# -*- coding: utf-8 -*-
"""
robot_writer.py (v3 - 最终修复版)
-------------------------------------------------
功能:
- v3修复: 新增全角标点到半角的自动映射，完美支持中文标点符号。
- v2修复: 全面支持中、英、数字及ASCII标点符号。
- v2修复: 修正了英文绘制函数中的bug，可正常书写英文。
"""
import pickle
import time
import math
from pymycobot.ultraArmP340 import ultraArmP340
from HersheyFonts import HersheyFonts

# ==================== 机械臂核心参数 ====================
COM_PORT = "COM35"      # 请根据实际情况修改
BAUDRATE = 115200
SPEED_MODE = 2

# ==================== 书写物理参数 ====================
PEN_Z_DOWN = -23
PEN_Z_UP = -19.5
PEN_Z_GAP = PEN_Z_UP + 3
SPEED_DRAW = 55
SPEED_MOVE = 55
PAUSE = 0.01

# ==================== 坐标系映射参数 ====================
ROBOT_START_X = 335.0
ROBOT_START_Y = -100.0

# ==================== 字体数据与偏移 ====================
try:
    with open("Chinese_strokes", "rb") as f:
        chinese_data = pickle.load(f)
except FileNotFoundError:
    print("❌ 错误：找不到中文字库文件 'Chinese_strokes'。")
    exit()

ASCII_OFFSET_X_FACTOR = -0.7
ASCII_OFFSET_Y_FACTOR = -0.23

# --- NEW: 全角转半角标点符号映射表 ---
PUNCTUATION_MAP = {
    '，': ',', '。': '.', '“': '"', '”': '"', '：': ':', 
    '；': ';', '？': '?', '！': '!', '（': '(', '）': ')',
    '《': '<', '》': '>', '‘': "'", '’': "'", '、': ','
}

# ==================== 底层绘制函数 (无改动) ====================

def draw_chinese_char(ua, ch, center_x, center_y, char_height_mm):
    if ch not in chinese_data: return
    strokes = chinese_data[ch]
    char_spacing_base = 70.0
    scale = char_height_mm / char_spacing_base
    for stroke in strokes:
        x0_raw = stroke[0]["y"] / 10.0
        y0_raw = stroke[0]["x"] / 10.0 - char_spacing_base
        px0 = center_x - x0_raw * scale
        py0 = center_y + y0_raw * scale
        ua.set_coords([px0, py0, PEN_Z_GAP], SPEED_MOVE)
        ua.set_coords([px0, py0, PEN_Z_DOWN], SPEED_DRAW)
        time.sleep(PAUSE)
        for pt in stroke[1:]:
            x_raw = pt["y"] / 10.0
            y_raw = pt["x"] / 10.0 - char_spacing_base
            px = center_x - x_raw * scale
            py = center_y + y_raw * scale
            ua.set_coords([px, py, PEN_Z_DOWN], SPEED_DRAW)
            time.sleep(PAUSE)
        ua.set_coords([px, py, PEN_Z_UP], SPEED_MOVE)
        time.sleep(PAUSE + 0.05)

def merge_segments_to_paths(segments):
    segs = [((x1, y1), (x2, y2)) for ((x1, y1), (x2, y2)) in segments]
    paths = []
    while segs:
        (start, end) = segs.pop(0)
        path = [start, end]
        extended = True
        while extended:
            extended = False
            for idx, ((x1, y1), (x2, y2)) in enumerate(segs):
                if (x1, y1) == path[-1]: path.append((x2, y2)); segs.pop(idx); extended = True; break
                elif (x2, y2) == path[-1]: path.append((x1, y1)); segs.pop(idx); extended = True; break
        extended = True
        while extended:
            extended = False
            for idx, ((x1, y1), (x2, y2)) in enumerate(segs):
                if (x2, y2) == path[0]: path.insert(0, (x1, y1)); segs.pop(idx); extended = True; break
                elif (x1, y1) == path[0]: path.insert(0, (x2, y2)); segs.pop(idx); extended = True; break
        paths.append(path)
    return paths

def draw_ascii_char(ua, ch, center_x, center_y, char_height_mm):
    hf = HersheyFonts(); hf.load_default_font("futural")
    ascii_unit_height = 100.0; hf.normalize_rendering(ascii_unit_height)
    ascii_scale = char_height_mm / ascii_unit_height
    raw_segments = hf.lines_for_text(ch)
    if not raw_segments: return
    paths = merge_segments_to_paths(raw_segments)
    if not paths: return
    all_x_raw = [pt[0] for path_item in paths for pt in path_item]
    if not all_x_raw: return
    center_offset_x_raw = (min(all_x_raw) + max(all_x_raw)) / 2.0
    ascii_offset_x = char_height_mm * ASCII_OFFSET_X_FACTOR
    ascii_offset_y = char_height_mm * ASCII_OFFSET_Y_FACTOR
    for path in paths:
        x_h0, y_h0 = path[0]
        x_adj_raw0 = x_h0 - center_offset_x_raw
        x_pen_offset0 = (y_h0 * ascii_scale) - (char_height_mm / 2.0)
        y_pen_offset0 = x_adj_raw0 * ascii_scale
        px0 = center_x + x_pen_offset0 + ascii_offset_x
        py0 = center_y + y_pen_offset0 + ascii_offset_y
        ua.set_coords([px0, py0, PEN_Z_GAP], SPEED_MOVE)
        ua.set_coords([px0, py0, PEN_Z_DOWN], SPEED_DRAW)
        time.sleep(PAUSE)
        for (x_h, y_h) in path[1:]:
            x_adj_raw = x_h - center_offset_x_raw
            x_pen_offset = (y_h * ascii_scale) - (char_height_mm / 2.0)
            y_pen_offset = x_adj_raw * ascii_scale
            px = center_x + x_pen_offset + ascii_offset_x
            py = center_y + y_pen_offset + ascii_offset_y
            ua.set_coords([px, py, PEN_Z_DOWN], SPEED_DRAW)
            time.sleep(PAUSE)
        ua.set_coords([px, py, PEN_Z_UP], SPEED_MOVE)
        time.sleep(PAUSE + 0.05)


# ==================== 高层控制函数 ====================

def draw_text_line(ua, text, a4_x_mm, a4_y_mm, char_height_mm, char_spacing_ratio):
    print(f"  正在书写: '{text}' at A4(x={a4_x_mm:.1f}, y={a4_y_mm:.1f})")
    current_a4_x_offset = 0

    for char in text:
        # --- MODIFIED: 核心修改逻辑 ---
        # 1. 检查原始字符是否是全角标点，如果是，则映射为半角用于绘制
        is_full_width_punct = char in PUNCTUATION_MAP
        char_to_draw = PUNCTUATION_MAP.get(char, char) # 如果是标点，则使用映射后的半角字符，否则用原字符

        # 2. 判断字符类型
        is_chinese = char_to_draw in chinese_data
        is_space = char_to_draw == ' '
        is_ascii_drawable = char_to_draw.isascii() and char_to_draw.isprintable() and not is_space and not is_chinese
        
        # 3. 决定字符宽度以保持布局。全角标点应占据一个汉字宽度
        char_width_mm = char_height_mm if (is_chinese or is_full_width_punct) else char_height_mm / 2.0
        
        if is_space:
            current_a4_x_offset += char_width_mm * char_spacing_ratio
            continue

        # 计算字符中心的A4坐标
        center_a4_x = a4_x_mm + current_a4_x_offset + char_width_mm / 2.0
        center_a4_y = a4_y_mm + char_height_mm / 2.0
        
        # 坐标转换
        robot_x = ROBOT_START_X - center_a4_y
        robot_y = ROBOT_START_Y + center_a4_x
        
        # 调用绘制函数
        if is_chinese:
            draw_chinese_char(ua, char_to_draw, robot_x, robot_y, char_height_mm)
        elif is_ascii_drawable:
            draw_ascii_char(ua, char_to_draw, robot_x, robot_y, char_height_mm)
        else:
            # 只有在所有条件都不满足时才跳过
            print(f"    (跳过不支持的字符: {char})")

        # 更新下一个字符的起始位置
        spacing_mm = char_width_mm * char_spacing_ratio
        current_a4_x_offset += spacing_mm


def execute_writing_session(ua, tasks):
    if not tasks:
        print("本部分没有书写任务。")
        return
    sorted_tasks = sorted(tasks, key=lambda t: (t['a4_y_mm'], t['a4_x_mm']))
    initial_task = sorted_tasks[0]
    initial_robot_x = ROBOT_START_X - initial_task['a4_y_mm']
    initial_robot_y = ROBOT_START_Y + initial_task['a4_x_mm']
    ua.set_coords([initial_robot_x, initial_robot_y, PEN_Z_GAP], SPEED_MOVE)
    time.sleep(1)
    for i, task in enumerate(sorted_tasks):
        print(f"  任务 {i+1}/{len(sorted_tasks)}...")
        draw_text_line(ua, task["text"], task["a4_x_mm"], task["a4_y_mm"], task["char_height_mm"], task["char_spacing_ratio"])
        time.sleep(0.2)


def execute_full_writing_process(tasks_part1, tasks_part2):
    ua = None
    try:
        print("🔧 初始化机械臂...")
        ua = ultraArmP340(COM_PORT, BAUDRATE)
        ua.set_speed_mode(SPEED_MODE)

        print("✅ 机械臂初始化完成。")

        if tasks_part1:
            print("\n" + "="*15 + " 开始书写第 1 部分 (上半页) " + "="*15)
            execute_writing_session(ua, tasks_part1)
            print("✅ 第 1 部分书写完成。")
        else:
            print("\nℹ️ 第 1 部分没有需要书写的任务。")
        
        print("\n↩️ 机械臂归零，准备进行第二部分书写...")
        ua.go_zero()
        print("‼️ 请手动向上移动纸张，将下一区域对准书写起始点。")
        print("   将在 5 秒后继续...")
        time.sleep(5)

        if tasks_part2:
            print("\n" + "="*15 + " 开始书写第 2 部分 (下半页) " + "="*15)
            execute_writing_session(ua, tasks_part2)
            print("✅ 第 2 部分书写完成。")
        else:
            print("\nℹ️ 第 2 部分没有需要书写的任务。")

        print("\n🎉🎉🎉 全部书写任务已完成！")
        ua.go_zero()

    except Exception as err:
        print(f"❌ 运行异常: {err}")
        if ua:
            try: ua.go_zero()
            except: pass
    finally:
        if ua:
            print("👋 程序结束，机械臂已归位。")